import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
from sklearn.metrics import silhouette_score
from torchvision import transforms
from model import UNet

def get_result(image_path, k, scale=0.2):
    # 读取并缩小图像
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    image = cv2.resize(image, (0, 0), fx=scale, fy=scale)  # 缩小图像为 0.2 倍
    original_shape = image.shape  # 保存缩放后图像的形状
    pixels = image.reshape((-1, 1))  # 展平为(像素数, 1)的形状
    pixels = np.float32(pixels)

    # 定义 KMeans 聚类
    kmeans = KMeans(n_clusters=k, n_init=20, random_state=42)
    kmeans.fit(pixels)  # 聚类

    # 重构图像
    labels = kmeans.labels_
    centers = np.uint8(kmeans.cluster_centers_)
    segmented_image = centers[labels.flatten()].reshape(original_shape)

    return segmented_image, pixels

def bisecting_kmeans(X, k):
    clusters = [X]
    labels = np.zeros(X.shape[0], dtype=int)
    cluster_indices = [np.arange(X.shape[0])]

    while len(clusters) < k:
        max_sse_cluster_idx = 0
        max_sse = -1
        for i, cluster in enumerate(clusters):
            if len(cluster) > 1:
                km = KMeans(n_clusters=1, n_init=20).fit(cluster)
                sse = km.inertia_
                if sse > max_sse:
                    max_sse = sse
                    max_sse_cluster_idx = i

        # 拆分最大的聚类
        cluster_to_split = clusters[max_sse_cluster_idx]
        km = KMeans(n_clusters=2, n_init=20).fit(cluster_to_split)
        split_labels = km.labels_

        indices_to_split = cluster_indices[max_sse_cluster_idx]
        clusters.pop(max_sse_cluster_idx)
        cluster_indices.pop(max_sse_cluster_idx)

        clusters.append(cluster_to_split[split_labels == 0])
        clusters.append(cluster_to_split[split_labels == 1])
        cluster_indices.append(indices_to_split[split_labels == 0])
        cluster_indices.append(indices_to_split[split_labels == 1])

        # 更新标签
        labels[indices_to_split[split_labels == 0]] = len(clusters) - 2
        labels[indices_to_split[split_labels == 1]] = len(clusters) - 1

    return labels

# 设定参数
image_path = 'DRIVE/test/image/01.png'
label_path = 'DRIVE/test/label/01.png'
k = 3 # 聚类数
scale_factor = 1  # 缩小比例

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), 0.5)
])

# 加载模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = UNet(in_channels=3, num_classes=1)
model.load_state_dict(torch.load('UNet.pth', map_location=device))
model.to(device)

# 预测并进行聚类
model.eval()
with torch.no_grad():
    img = Image.open(image_path)
    label_img = Image.open(label_path)

    # 缩小图像
    img = img.resize((int(img.width * scale_factor), int(img.height * scale_factor)))
    label_img = label_img.resize((int(label_img.width * scale_factor), int(label_img.height * scale_factor)))

    img = transform(img)
    label_img = transform(label_img)
    img = torch.unsqueeze(img, dim=0)

    pred = model(img.to(device))
    pred_numpy = torch.sigmoid(pred).squeeze().cpu().numpy() * 255
    batch_size, channels, height, width = pred.size()
    outputs_flat = pred.view(batch_size, -1).cpu().numpy()

    output_map = outputs_flat[0].reshape(-1, 1)

    # KMeans 聚类
    kmeans = KMeans(n_clusters=k, n_init=20, random_state=42).fit(output_map)
    clustered_map_kmeans = kmeans.labels_.reshape(height, width)

    # KMeans++ 聚类
    kmeans_plus = KMeans(n_clusters=k, init='k-means++', n_init=20, random_state=42).fit(output_map)
    clustered_map_kmeans_plus = kmeans_plus.labels_.reshape(height, width)

    # Bisecting KMeans 聚类
    clustered_map_bisecting = bisecting_kmeans(output_map, k)

    # 无深度学习的 KMeans 聚类
    image_no_deep, pixels_nodeep = get_result(image_path, k, scale_factor)


    # 分别绘制每个结果
    plt.figure(figsize=(6, 6))
    plt.title('Label Image')
    plt.imshow(label_img.squeeze().cpu().numpy(), cmap='gray')
    plt.axis('off')
    plt.show()

    plt.figure(figsize=(6, 6))
    plt.title('KMeans Clustered Image')
    plt.imshow(clustered_map_kmeans, cmap='gray')
    plt.axis('off')
    plt.show()

    plt.figure(figsize=(6, 6))
    plt.title('KMeans++ Clustered Image')
    plt.imshow(clustered_map_kmeans_plus, cmap='gray')
    plt.axis('off')
    plt.show()

    plt.figure(figsize=(6, 6))
    plt.title('Bisecting KMeans')
    plt.imshow(clustered_map_bisecting.reshape(height, width), cmap='gray')
    plt.axis('off')
    plt.show()

    plt.figure(figsize=(6, 6))
    plt.title('KMeans Clustered (No Deep)')
    plt.imshow(image_no_deep, cmap='gray')
    plt.axis('off')
    plt.show()

    # 计算 Silhouette 系数
    silhouette_kmeans = silhouette_score(output_map, kmeans.labels_)
    silhouette_kmeans_plus = silhouette_score(output_map, kmeans_plus.labels_)
    silhouette_bisecting = silhouette_score(output_map, clustered_map_bisecting)
    silhouette_no_deep = silhouette_score(pixels_nodeep, image_no_deep.reshape(-1, 1))

    print(f'Silhouette Coefficient for KMeans: {silhouette_kmeans:.4f}')
    print(f'Silhouette Coefficient for KMeans++: {silhouette_kmeans_plus:.4f}')
    print(f'Silhouette Coefficient for Bisecting KMeans: {silhouette_bisecting:.4f}')
    print(f'Silhouette Coefficient for KMeans (No Deep): {silhouette_no_deep:.4f}')